import os
import sys
from typing import List, Dict, Any

from langchain.document_loaders import PyPDFLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.vectorstores import Chroma
from langchain.embeddings import HuggingFaceEmbeddings
from langchain.chains import RetrievalQA
from langchain_community.llms import DeepSeek


def load_pdf_documents(pdf_path: str):
    """从指定路径加载PDF文档"""
    print(f"正在加载PDF文件: {pdf_path}")
    loader = PyPDFLoader(pdf_path)
    documents = loader.load()
    print(f"成功加载文档，共 {len(documents)} 页")
    return documents


def split_documents(documents, chunk_size=1000, chunk_overlap=200):
    """将文档分割成更小的块"""
    text_splitter = RecursiveCharacterTextSplitter(
        chunk_size=chunk_size,
        chunk_overlap=chunk_overlap,
        length_function=len,
    )
    chunks = text_splitter.split_documents(documents)
    print(f"文档已分割为 {len(chunks)} 个块")
    return chunks


def create_vector_store(chunks):
    """创建向量存储"""
    # 使用HuggingFace的嵌入模型
    embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
    
    # 创建Chroma向量存储
    vector_store = Chroma.from_documents(documents=chunks, embedding=embeddings)
    print("向量存储创建完成")
    return vector_store


def setup_qa_chain(vector_store):
    """设置问答链"""
    # 使用DeepSeek-v3模型
    llm = DeepSeek(model="deepseek-ai/deepseek-v3")
    
    # 创建检索QA链
    qa_chain = RetrievalQA.from_chain_type(
        llm=llm,
        chain_type="stuff",
        retriever=vector_store.as_retriever(search_kwargs={"k": 3}),
        return_source_documents=True,
    )
    return qa_chain


def ask_question(qa_chain, question: str):
    """向RAG系统提问"""
    print(f"\n问题: {question}")
    result = qa_chain({"query": question})
    print("\n回答:")
    print(result["result"])
    print("\n参考来源:")
    for i, doc in enumerate(result["source_documents"]):
        print(f"来源 {i+1}:\n{doc.page_content}\n")
    return result


def main():
    # 检查命令行参数
    if len(sys.argv) < 2:
        print("使用方法: python main.py <pdf文件路径>")
        sys.exit(1)
    
    pdf_path = sys.argv[1]
    if not os.path.exists(pdf_path):
        print(f"错误: 文件 {pdf_path} 不存在")
        sys.exit(1)
    
    # 加载文档
    documents = load_pdf_documents(pdf_path)
    
    # 分割文档
    chunks = split_documents(documents)
    
    # 创建向量存储
    vector_store = create_vector_store(chunks)
    
    # 设置QA链
    qa_chain = setup_qa_chain(vector_store)
    
    # 交互式问答
    print("\n初始化完成! 您可以开始提问 (输入'退出'结束)")
    while True:
        question = input("\n请输入您的问题: ")
        if question.lower() in ["退出", "exit", "quit"]:
            break
        ask_question(qa_chain, question)


if __name__ == "__main__":
    main()